# 这段代码的主要功能是实现一个自定义的图像数据加载与增强系统，特别用于无人机(UAV)图像和基于 GPS 信息的地理定位任务。具体模块功能如下：
# 1. 数据处理模块
# Dataset_gallery 类
# 主要负责加载带有 GPS 信息的图像数据。
# 数据格式通过一个文件指定，包含图像路径和 GPS 坐标。
# 支持通过 filter_mode 参数过滤特定类型的文件。
# Dataset_query 类
# 加载查询数据集，从指定目录中读取图像，并结合一个 GPS 信息文件获取对应的位置信息。
# 支持处理带层级目录结构的文件路径。
# 返回图像、对应的 GPS 坐标，以及路径。
# test_collate_fn 函数
# 将批量图像数据整合为张量，同时返回经纬度信息和图像路径。
# 用于数据加载器中的批处理阶段。
# 2. 数据增强模块
# Query_transforms
# 对图像进行镜像填充和裁剪操作。
# 保证图像尺寸一致，为模型输入准备数据。
# CenterCrop
# 实现图像的中心裁剪，使图像保持方形。
# 用于去除多余边缘信息。
# RotateAndCrop
# 随机旋转图像，并根据旋转后的角度裁剪出一个新的图像。
# 模拟无人机在不同视角下采集的数据。
# RandomCrop
# 随机裁剪图像的中心部分，保持比例一致。
# 用于模拟不同采样位置的图像。
# RandomErasing
# 随机擦除图像的一部分，用于增强模型对遮挡或信息丢失的鲁棒性。
# 3. 测试模块
# 主程序中：
# 使用 RandomErasing、RandomAffine 和 RotateAndCrop 等变换方法对图像进行增强和可视化保存。
# 生成增强后图像并标记其中心，用于验证图像变换效果

from torch.utils.data import Dataset
import numpy as np
import cv2
from PIL import Image
from math import cos, sin, pi
import random
import math
import glob
import os
import torch


def test_collate_fn(batch):
    """
    合并一个 batch 的数据。
    输入：list，每个元素是 Dataset 类 __getitem__ 方法返回的结果 (image, infos, path)。
    输出：
        - 拼接的图像张量
        - 经纬度信息（二维数组）
        - 图像路径列表
    """
    image, infos, path = zip(*batch)  # 解压批次数据
    E, N = zip(*infos)  # 分离纬度 (E) 和经度 (N)
    E = np.array(E)  # 转为 NumPy 数组
    N = np.array(N)
    infos = np.stack((E, N), axis=1)  # 合并纬度和经度信息为二维数组
    return torch.stack(image, dim=0), np.stack(infos, axis=0), path  # 返回图像张量、信息数组和路径

class Dataset_gallery(Dataset):
    def __init__(self, filename, transformer, filter_mode=None):
        """
        加载包含图像路径和位置信息的文件。
        参数：
            - filename：包含图像路径和 GPS 信息的文件路径
            - transformer：图像增强变换
            - filter_mode：过滤模式，若指定，则只加载路径中包含此字符串的文件
        """
        super(Dataset_gallery, self).__init__()
        self.filename = filename  # 数据文件路径
        self.transformer = transformer  # 图像变换
        info_list = []  # 用于存储路径和 GPS 信息
        with open(self.filename, "r") as F:
            context = F.readlines()
            for line in context:
                info = line.strip().split(" ")  # 分割路径和 GPS 信息
                if filter_mode is not None and filter_mode not in info[0]:
                    continue  # 跳过不匹配的文件
                if info[0].endswith(".jpg"):  # 仅处理 .jpg 文件
                    info_list.append(info)
        self.info_list = info_list  # 保存所有符合条件的信息

    def __getitem__(self, ind):
        """
        获取单个数据项。
        返回：
            - 图像（经过增强）
            - GPS 信息（纬度、经度）
            - 图像路径
        """
        info = self.info_list[ind]
        image = Image.open(info[0])  # 加载图像
        image = self.transformer(image)  # 应用增强变换
        return image, [float(info[1]), float(info[2])], [info[0]]

    def __len__(self):
        """返回数据集大小（总样本数）。"""
        return len(self.info_list)
    

class Dataset_query(Dataset):
    def __init__(self, query_dir, transformer, info_file="/path/to/info.txt"):
        """
        加载查询数据集，结合文件夹中的图像路径和 GPS 信息。
        参数：
            - query_dir：查询数据的根目录
            - transformer：图像增强变换
            - info_file：包含 GPS 信息的文件路径
        """
        super(Dataset_query, self).__init__()
        self.query_dir = query_dir
        self.transformer = transformer
        self.query_images = glob.glob(os.path.join(query_dir, "*/*.JPG"))  # 查找所有 .JPG 文件

        # 读取 GPS 信息，保存为字典
        configDict = {}
        with open(info_file, "r") as F:
            context = F.readlines()
            for line in context:
                splitLineList = line.split(" ")  # 分割行数据
                configDict[splitLineList[0].split("/")[-2]] = [
                    float(splitLineList[1].split("E")[-1]),
                    float(splitLineList[2].split("N")[-1]),
                    splitLineList[0]
                ]
        self.configDict = configDict  # 保存字典

    def __getitem__(self, ind):
        """
        获取单个数据项。
        返回：
            - 图像（经过增强）
            - GPS 信息（纬度、经度）
            - 图像路径
        """
        image_path = self.query_images[ind]
        name = image_path.split("/")[-2]  # 提取路径中的文件夹名
        info = self.configDict[name]  # 获取对应的 GPS 信息
        image = Image.open(image_path)  # 加载图像
        image = self.transformer(image)  # 应用增强变换
        return image, info[:2], info[2]

    def __len__(self):
        """返回数据集大小（总样本数）。"""
        return len(self.query_images)


class Query_transforms(object):
    def __init__(self, pad=20, size=256):
        """
        初始化增强方法。
        参数：
            - pad：边缘填充宽度
            - size：最终裁剪尺寸
        """
        self.pad = pad
        self.size = size

    def __call__(self, img):
        """
        对图像执行增强操作。
        1. 镜像填充
        2. 裁剪到指定尺寸
        """
        img_ = np.array(img).copy()
        img_part = img_[:, 0:self.pad, :]  # 提取填充区域
        img_flip = cv2.flip(img_part, 1)  # 水平翻转填充区域
        image = np.concatenate((img_flip, img_), axis=1)  # 拼接原图和填充区域
        image = image[:, 0:self.size, :]  # 裁剪到指定尺寸
        image = Image.fromarray(image.astype('uint8')).convert('RGB')  # 转换为 PIL 格式
        return image


class CenterCrop(object):
    """Composes several transforms together.

    Args:
        transforms (list of ``Transform`` objects): list of transforms to compose.

    Example:
        >>> transforms.Compose([
        >>>     transforms.CenterCrop(10),
        >>>     transforms.ToTensor(),
        >>> ])
    """

    def __init__(self):
        pass

    def __call__(self, img):
        img_ = np.array(img).copy()
        h, w, c = img_.shape
        min_edge = min((h, w))
        if min_edge == h:
            edge_lenth = int((w - min_edge) / 2)
            new_image = img_[:, edge_lenth:w - edge_lenth, :]
        else:
            edge_lenth = int((h - min_edge) / 2)
            new_image = img_[edge_lenth:h - edge_lenth, :, :]
        assert new_image.shape[0] == new_image.shape[1], "the shape is not correct"
        cv2.imshow("query", cv2.resize(new_image, (512, 512)))
        image = Image.fromarray(new_image.astype('uint8')).convert('RGB')
        return image


class RotateAndCrop(object):
    def __init__(self, rate, output_size=(512, 512), rotate_range=360):
        self.rate = rate
        self.output_size = output_size
        self.rotate_range = rotate_range

    def __call__(self, img):
        img_ = np.array(img).copy()

        def getPosByAngle(img, angle):
            h, w, c = img.shape
            y_center = h // 2
            x_center = w//2
            r = h // 2
            angle_lt = angle - 45
            angle_rt = angle + 45
            angle_lb = angle + 135
            angle_rb = angle + 225
            angleList = [angle_lt, angle_rt, angle_lb, angle_rb]
            pointsList = []
            for angle in angleList:
                x1 = x_center + r * cos(angle * pi / 180)
                y1 = y_center + r * sin(angle * pi / 180)
                pointsList.append([x1, y1])
            pointsOri = np.float32(pointsList)
            pointsListAfter = np.float32(
                [[0, 0], [0, self.output_size[0]], [self.output_size[0], self.output_size[1]], [self.output_size[1], 0]])
            M = cv2.getPerspectiveTransform(pointsOri, pointsListAfter)
            res = cv2.warpPerspective(
                img, M, (self.output_size[0], self.output_size[1]))
            return res

        if np.random.random() > self.rate:
            image = img
        else:
            angle = int(np.random.random()*self.rotate_range)
            new_image = getPosByAngle(img_, angle)
            image = Image.fromarray(new_image.astype('uint8')).convert('RGB')
        return image


class RandomCrop(object):
    def __init__(self, rate=0.2):
        self.rate = rate

    def __call__(self, img):
        img_ = np.array(img).copy()
        h, w, c = img_.shape
        random_width = int(np.random.random()*self.rate*w)
        random_height = int(np.random.random()*self.rate*h)
        x_l = random_width
        x_r = w-random_width
        y_l = random_height
        y_r = h-random_height
        new_image = img_[y_l:y_r, x_l:x_r, :]
        image = Image.fromarray(new_image.astype('uint8')).convert('RGB')
        return image


class RandomErasing(object):
    """ Randomly selects a rectangle region in an image and erases its pixels.
        'Random Erasing Data Augmentation' by Zhong et al.
        See https://arxiv.org/pdf/1708.04896.pdf
    Args:
         probability: The probability that the Random Erasing operation will be performed.
         sl: Minimum proportion of erased area against input image.
         sh: Maximum proportion of erased area against input image.
         r1: Minimum aspect ratio of erased area.
         mean: Erasing value. 
    """

    def __init__(self, probability=0.5, sl=0.02, sh=0.3, r1=0.3):
        self.probability = probability
        self.sl = sl
        self.sh = sh
        self.r1 = r1

    def __call__(self, img):
        if random.uniform(0, 1) > self.probability:
            return img

        img_ = np.array(img).copy()

        mean = np.mean(np.mean(img_, 0), 1)

        for _ in range(100):
            area = img_.shape[0] * img_.shape[1]

            target_area = random.uniform(self.sl, self.sh) * area
            aspect_ratio = random.uniform(self.r1, 1/self.r1)

            h = int(round(math.sqrt(target_area * aspect_ratio)))
            w = int(round(math.sqrt(target_area / aspect_ratio)))

            if w < img_.shape[1] and h < img_.shape[0]:
                x1 = random.randint(0, img_.shape[1] - h)
                y1 = random.randint(0, img_.shape[0] - w)

                img_[x1:x1+h, y1:y1+w, 0] = mean[0]
                img_[x1:x1+h, y1:y1+w, 1] = mean[1]
                img_[x1:x1+h, y1:y1+w, 2] = mean[2]
                img = Image.fromarray(img_.astype('uint8')).convert('RGB')
                return img

        return img


if __name__ == '__main__':
    image_path = "/home/ming.dai/workspace/dataset/DenseUAV/data_2022/test/drone/002810/H80.JPG"
    image = Image.open(image_path)

    re = RandomErasing(probability=1.0)

    from torchvision import transforms
    ra = transforms.RandomAffine(180)

    rac = RotateAndCrop(rate=1.0)

    target_dir = "/home/ming.dai/workspace/code/DenseUAV/visualization/rotate_crop"
    import os
    os.makedirs(target_dir, exist_ok=True)
    for ind in range(10):
        image_ = rac(image)
        image_ = np.array(image_)
        image_ = image_[:, :, [2, 1, 0]]
        h, w = image_.shape[:2]
        image_ = cv2.circle(
            image_.copy(), (int(w/2), int(h/2)), 3, (0, 0, 255), 2)
        cv2.imwrite(os.path.join(target_dir, "{}.jpg".format(ind)), image_)
